from random import gauss
from statistics import variance
import numpy as np
from scipy.cluster.hierarchy import average
from scipy.stats import norm


def gauss_pdf(x, mu, sigma):
    """计算高斯概率密度函数"""
    return norm.pdf(x, mu, sigma)


def averageWeight(data, weights, total_weight):
    """计算加权平均值"""
    if total_weight == 0:
        return 0
    return sum(data[i] * weights[i] for i in range(len(data))) / total_weight


def varianceWeight(data, weights, mean, total_weight):
    """计算加权方差"""
    if total_weight == 0:
        return 1
    variance = sum(weights[i] * (data[i] - mean) ** 2 for i in range(len(data))) / total_weight
    return max(variance, 0.1)  # 防止方差为0


def isSame(cur, now, tolerance=1e-6):
    """判断两个参数集合是否相同（在容忍度范围内）"""
    return all(abs(a - b) < tolerance for a, b in zip(cur, now))


def calcEM(height):
    N = len(height)
    gp = 0.5  # 女孩的概率
    bp = 0.5  # 男孩的概率
    gmu, gsigma = min(height), 1  # 女孩身高的均值和标准差
    bmu, bsigma = max(height), 1  # 男孩身高的均值和标准差
    ggamma = [0] * N  # 属于女孩的后验概率
    bgamma = [0] * N  # 属于男孩的后验概率

    cur = [gp, bp, gmu, gsigma, bmu, bsigma]
    now = []

    times = 0
    while times < 100:
        # E步：计算后验概率
        i = 0
        for x in height:
            ggamma[i] = gp * gauss_pdf(x, gmu, gsigma)
            bgamma[i] = bp * gauss_pdf(x, bmu, bsigma)
            s = ggamma[i] + bgamma[i]
            if s > 0:  # 防止除零
                ggamma[i] /= s
                bgamma[i] /= s
            i += 1

        # M步：更新参数
        gn = sum(ggamma)  # 女孩的期望数量
        bn = sum(bgamma)  # 男孩的期望数量

        gp = float(gn) / float(N)  # 更新女孩概率
        bp = float(bn) / float(N)  # 更新男孩概率

        gmu = averageWeight(height, ggamma, gn)
        gsigma = np.sqrt(varianceWeight(height, ggamma, gmu, gn))
        bmu = averageWeight(height, bgamma, bn)
        bsigma = np.sqrt(varianceWeight(height, bgamma, bmu, bn))

        now = [gp, bp, gmu, gsigma, bmu, bsigma]

        # 检查收敛
        if isSame(cur, now):
            break

        cur = now
        print("Times:\t", times)
        print("Girl mean/gsigma:\t", gmu, gsigma)
        print("Boy mean/bsigma:\t", bmu, bsigma)
        print("Boy/Girl proportion:\t", bp, gp)
        print("Boy/Girl count:\t", bn, gn, bn + gn)
        print("\n")
        times += 1

    return now


# 测试代码
if __name__ == "__main__":
    # 生成测试数据：女孩身高~N(160, 5)，男孩身高~N(175, 6)
    np.random.seed(42)#设置随机数生成器的初始状态，确保后续随机操作可重复。
    girl_heights = np.random.normal(160, 5, 100)#均值160，标准差5，生成100个数
    boy_heights = np.random.normal(175, 6, 100)
    heights = np.concatenate([girl_heights, boy_heights])#拼接两个数组

    result = calcEM(heights)
    print("Final result:", result)